#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   myopiadataloader.py
@Time    :   2020/06/21 14:02:40
@Author  :   caotian6666 
@Version :   1.0
@Contact :   caotiandyx@163.com
@License :   (C)Copyright 2017-2018, Liugroup-NLPR-CASIA
@Desc    :   None
'''

import os
import sys
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import paddle
import paddle.fluid as fluid
import cv2
import random
import numpy as np

def transform_img(img):
    img=cv2.resize(img,(224,224))
    img=np.transpose(img,(2,0,1))
    img=img.astype('float32')
    img=img/255.
    img=img * 2.0-1.0
    return  img
def data_loader(datadir,batch_size=10,mode='train'):
    filenames=os.listdir(datadir)
    def reader():
        if mode=='train':
            random.shuffle(filenames)
        batch_imgs=[]
        batch_labels=[]
        for name in filenames:
            filepath=os.path.join(datadir,name)
            img=cv2.imread(filepath)
            img=transform_img(img)
            if name[0]=='H' or name[0] == 'N':
                label=0
            elif name[0]=='P':
                label=1
            else:
                raise('Not excepted file name')
            batch_imgs.append(img)
            batch_labels.append(label)
            if len(batch_imgs)==batch_size:
                imgs_array=np.array(batch_imgs).astype('float32')
                labels_array=np.array(batch_labels).astype('float32').reshape(-1,1)
                yield imgs_array,labels_array
                batch_imgs=[]
                batch_labels=[]
        if len(batch_imgs) > 0:
            imgs_array=np.array(batch_imgs).astype('float32')
            labels_array=np.array(batch_labels).astype('float32')
            yield imgs_array,labels_array
    return reader
def valid_data_loader(datadir,csvfile,batch_size=10,mode='valid'):
    filelists=open(csvfile).readlines()
    def reader():
        batch_imgs=[]
        batch_labels=[]
        for line in filelists[1:]:
            line=line.strip().split(',')
            name=line[1]
            label=int(line[2])
            filepath=os.path.join(datadir,name)
            img=cv2.imread(filepath)
            img=transform_img(img)
            batch_imgs.append(img)
            batch_labels.append(label)

            if len(batch_imgs)==batch_size:
                imgs_array=np.array(batch_imgs).astype('float32')
                labels_array=np.array(batch_labels).astype('float32').reshape(-1,1)
                yield imgs_array,labels_array
                batch_imgs=[]
                batch_labels=[]
        if len(batch_imgs) > 0:
            imgs_array=np.array(batch_imgs).astype('float32')
            labels_array=np.array(batch_labels).astype('float32').reshape(-1,1)
            yield imgs_array,labels_array
    return reader

if __name__ == '__main__':
    datadir='./data/PALM-Training400/PALM-Training400'
    train_loader=data_loader(datadir,batch_size=10,mode='train')
    data_reader=train_loader()
    data=next(data_reader)
    print("data[0] shape is:{},data[1] shape is:{}".format(data[0].shape,data[1].shape))

